fix: serialize token refresh across processes and fix OpenCode's dynamic Authorization plugin (#190)#197
Conversation
There was a problem hiding this comment.
Pull request overview
This PR addresses intermittent Databricks OAuth token invalidation in concurrent ucode opencode / ucode gemini sessions by (1) serializing databricks auth token --force-refresh across processes and coalescing redundant refreshes, and (2) fixing OpenCode’s long-lived token staleness by generating an OpenCode chat.headers plugin that refreshes Authorization per request.
Changes:
- Add a cross-process
flock+ sentinel-based coalescing around--force-refreshinget_databricks_token. - Generate and write a ucode-managed OpenCode plugin that sets a fresh
Authorizationheader per request (scoped to configured Databricks provider IDs). - Extend
ucode revertto delete the generated OpenCode plugin file, and add tests for locking/coalescing and plugin runtime behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
src/ucode/databricks.py |
Adds cross-process locking + coalescing sentinel around --force-refresh to avoid refresh-token reuse races. |
src/ucode/agents/opencode.py |
Generates/writes an OpenCode chat.headers plugin to refresh Authorization per request with TTL caching. |
src/ucode/cli.py |
Updates revert to delete the generated OpenCode plugin file. |
tests/test_databricks.py |
Adds coverage for lock-path scoping, coalescing behavior, and sentinel-touch semantics. |
tests/test_agent_opencode.py |
Adds coverage for plugin generation and a Node-executed runtime behavior test for hook input shape. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…lock (databricks#190) Concurrent long-lived ucode sessions each ran `databricks auth token --force-refresh` independently, racing on the shared OAuth cache and redeeming (and revoking) the rotating refresh token concurrently. Serialize --force-refresh calls per workspace+profile behind a cross-process flock under APP_DIR, and coalesce redundant refreshes: if a peer force-refreshed within TOKEN_REFRESH_COALESCE_SECONDS (60s), drop --force-refresh and reuse the now-current cached token instead of redeeming the refresh token again. Degrades gracefully to no locking when fcntl is unavailable (Windows) or lock acquisition fails for any reason -- a stuck/missing lock must never block a token refresh. The non-force-refresh (cheap read) path is unchanged and unlocked.
…s plugin (databricks#190) opencode reads provider.options.headers.Authorization from opencode.json once at startup and never again, so a long-lived session keeps sending a bearer token that may since have been rotated or revoked. Ship a ucode-managed opencode plugin (chat.headers hook, auto-loaded from the opencode plugin dir) that sets a fresh Authorization header on every chat request for our Databricks providers, sourcing the token from `ucode auth-token` but cached in-process with a short TTL (60s) plus an in-flight promise guard so it doesn't shell out on every request or spawn concurrent subprocesses. The static apiKey/headers in opencode.json are left untouched as the bootstrap/fallback the plugin fails open to on any refresh error. write_tool_config now also writes the plugin, scoped to only the provider ids actually configured. `ucode revert` deletes the (entirely ucode-generated, no-backup-needed) plugin file.
…ks#190) The generated chat.headers plugin guarded provider scoping with input.provider?.info?.id, but opencode's real hook input.provider is the runtime Provider record itself (id/name/env/options/source/models) -- id lives directly on input.provider, never under a nested .info. The guard's !providerId check was therefore always true, so the hook returned immediately on every request and never refreshed Authorization for any provider. Axis B was a complete no-op in production despite passing all prior tests, because those tests only asserted on the generated JS's string content, never its behavior against a real hook invocation. Found via an isolated end-to-end test against the real opencode 1.17.10 binary (separate $HOME/XDG dirs, real Databricks auth, no contact with the live ucode install): a well-formed-but-expired JWT baked into the static provider.options.headers.Authorization/apiKey still produced 'Bad Request: Invalid Token' with the plugin loaded and firing -- diagnostic instrumentation showed the guard's providerId was always undefined. Fixed to input.provider?.id and reverified against the same isolated harness: the plugin now fetches a fresh token and the request succeeds despite the stale static config. Adds TestRenderAuthPluginRuntimeBehavior, which shells out to a real node process and invokes the generated plugin's chat.headers hook against the actual runtime input shape (not a mocked one) -- verified to fail against the old input.provider?.info?.id line and pass against the fix, so this class of bug is now covered by an assertion that actually exercises behavior rather than string content.
…headers (review) Addresses two Copilot review comments on databricks#197: - fetchToken() returned stdout.trim() unconditionally. If `ucode auth-token` ever exits 0 with empty/whitespace-only stdout, the plugin would cache and send Authorization: Bearer <empty>, clobbering a possibly-still-valid bootstrap token for the full TTL window instead of failing open. Empty output now throws, so the hook's catch block fails open as intended. - The chat.headers hook assumed output.headers always exists. If opencode ever invokes the hook without a pre-populated headers object, the assignment would throw -- caught silently by the existing try/catch, but the refresh would never apply. Defensively initialize output.headers before writing to it. Both are proven by new Node-executed runtime tests (matching the class of test that caught the input.provider?.id bug): each new test was verified to fail against the pre-fix code and pass against the fix, not just pass trivially.
d2c3638 to
d259934
Compare
|
Addressed both review comments in
Both are covered by new Node-executed runtime tests ( Also rebased onto latest |
…ss (review) Addresses two more Copilot review comments on databricks#197 (both in the TestRenderAuthPluginRuntimeBehavior test harness, not production code): - _run_plugin() hard-coded `python3` when patching build_auth_token_argv for the fake ucode binary. python3 isn't guaranteed to be on PATH in every environment. Use sys.executable so the fake script runs under the same interpreter as the test. - The Node harness imported the generated plugin via plugin_path.as_posix(), a bare filesystem path. That's not a valid ESM import specifier on Windows (C:\...). Use Path.as_uri() for a portable file:// URL. The other two comments on this PR (opencode.py lines 233/275) are stale duplicates of the two already fixed in d259934 -- same comment IDs, diff_hunk frozen at review-creation time before that fix landed. No action needed there beyond the earlier reply.
|
Two more comments came in after the rebase (
The other two comments on this thread (opencode.py lines 233/275) are stale duplicates of the ones already fixed in |
…evert errors (review) Addresses three more Copilot review comments on databricks#197: - _touch_sentinel() could raise (swallowed) FileNotFoundError when its parent directory doesn't exist yet -- true whenever fcntl is unavailable (Windows), since _refresh_lock returns before ever creating APP_DIR on that branch. Silently disabled coalescing on every such platform. Now creates its own parent directory. - _refresh_lock()'s acquisition loop retried on *any* OSError from flock(), not just contention. An unexpected error (EBADF, or a filesystem that doesn't support flock at all) would spin for the full _LOCK_ACQUIRE_TIMEOUT_SECONDS (30s) before falling back to unlocked, instead of failing back immediately. Now only retries on BlockingIOError (the actual contention case) and bails out right away on anything else. - _remove_opencode_auth_plugin() silently swallowed OSError, so a real removal failure (permissions, filesystem issues) would leave `ucode revert` reporting "unchanged" and proceeding as if nothing were wrong. Now raises RuntimeError, matching restore_file's existing convention elsewhere in revert. All three are covered by new tests, each verified to fail against the pre-fix code (including the 30s-spin one, confirmed to actually take ~30s before the fix and <2s after) and pass against the fix.
|
Addressed the 3 new comments in
All three have new tests, each verified to fail against the pre-fix code and pass against the fix. |
Summary
Fixes #190. Running multiple concurrent
ucode opencodesessions against the same workspace/profile intermittently producesForbidden: Invalid Tokenon all sessions, and a single long-lived session eventually hits the same error after its token outlives the process's in-memory startup config. Two independent root causes, fixed separately:Axis A — concurrent
--force-refreshraces on the shared Databricks OAuth cache. Everyucode opencode/ucode geminiprocess runs its own background timer that callsdatabricks auth token --force-refreshevery 30 minutes, with no locking anywhere. When two sessions' timers land close together, both redeem the same rotating refresh token concurrently; the loser presents an already-rotated refresh token, tripping the OAuth server's reuse detection and revoking the whole token family — invalidating the access token that was just issued.Axis B — OpenCode reads a static token once at startup and never again. OpenCode has no
apiKeyHelperequivalent (unlike Claude Code) and only resolves{env:}/{file:}config placeholders at process start, so a long-running session keeps sending whatever token it read on launch even after ucode's background thread rewrites the file with a fresh one.Changes
databricks.py— wrap--force-refreshin a cross-processflock(per workspace+profile lockfile under~/.ucode), with a coalescing sentinel: if a peer refreshed within the last 60s, drop--force-refreshand reuse the cached token instead of redeeming again. Degrades gracefully to no-lock on Windows (nofcntl) or onOSError. The non-force read path is unchanged. The sentinel is marked on attempt, not success — Databricks rotates the refresh token on the redemption attempt reaching the server, not on our ability to parse a token out of the response, so marking only on success would leave a coalescing gap on a slow-but-successful peer.agents/opencode.py— ship a ucode-managed OpenCode plugin (chat.headershook, auto-loaded from the plugin dir) that sets a freshAuthorizationheader on every chat request for our Databricks providers, sourced fromucode auth-tokenwith an in-memory TTL cache (60s) + an in-flight promise guard so it doesn't shell out on every single request or spawn concurrent subprocesses. The staticapiKey/headersinopencode.jsonare left as-is as a bootstrap/fallback the plugin overrides live and fails open to on any refresh error.write_tool_confignow also writes the plugin, scoped to the provider ids actually configured.cli.py—ucode revertdeletes the (fully-generated, no backup needed) plugin file.A note on how Axis B's bug was actually found
The generated plugin's provider-scoping guard originally read
input.provider?.info?.id. Every unit test I wrote for it passed, because those tests only asserted on the generated JS's string content — one test literally asserted"input.provider?.info?.id" in jsas the expected string. Against a realopencodeprocess,input.provideris the runtimeProviderrecord itself (id/name/env/options/source/models) —idlives directly on it, not nested under.info. The guard's!providerIdcheck was therefore always true, so the hook returned immediately on every request and never refreshed anything — Axis B was a complete no-op in production.This only surfaced once I ran the plugin end-to-end against the real opencode 1.17.10 binary in an isolated environment (separate
$HOME/XDG dirs, real Databricks auth, no contact with any other live install) with a deliberately expired-but-well-formed JWT baked into the static config. It still failed withBad Request: Invalid Tokenwith the plugin loaded; diagnostic instrumentation showed the guard'sproviderIdwas alwaysundefined. Fixed toinput.provider?.idand reverified against the same harness — the plugin now fetches a fresh token and the request succeeds despite the stale static config.Added
TestRenderAuthPluginRuntimeBehavior, which shells out to a realnodeprocess and invokes the generated plugin'schat.headershook against the actual runtime input shape (skips ifnodeisn't onPATH). Verified it fails against the oldinput.provider?.info?.idline and passes against the fix, so this class of bug — assumed hook shape vs. real hook shape — now has coverage that isn't just a string-content assertion.Testing
uv run pytest: 881 passed, 37 skipped, 1 pre-existing unrelated failure (test_e2e_user_agent.py::TestClaudeUserAgent::test_user_agent_arrives_at_gateway, a real-claude-binary e2e test that fails identically onmain— verified viagit stash).uv run ruff check ./uv run ruff format --check/ty check src/: all clean.uv tool install --reinstall .) against a real Databricks workspace, plus an isolated end-to-end harness (separate$HOME, real auth, no contact with any live install) used to find and confirm the Axis B fix above.ucode opencodesessions run for ~18 hours against a real workspace, crossing the ~60-minute token lifetime boundary roughly 17+ times each. ZeroForbidden/Invalid Token/Unauthorized/Bad Requestauth failures in that entire window (verified via log inspection), with the background refresh cycling cleanly across all three sessions without collision.Follow-ups considered out of scope for this PR
ucode revert's plugin-file deletion doesn't check dry-run mode (only relevant ifrevertever runs under--dry-run).